#!/usr/bin/env python3
"""
Preprocess the Infinigen dataset to mimic the output of preprocess_hypersim.py.

This script reads data from the original pre-rendered Infinigen directory
structure, which is {scene}/{data_type}/{camera}. It applies a percentile-based
tonemapping method to the HDR .exr images, resizes the image and depth map to
a target resolution, adjusts the camera intrinsics accordingly, and saves the
processed outputs in a new directory structure with filenames matching the
Hypersim preprocessing script's output.

Usage:
    python preprocess_infinigen.py --infinigen_dir /path/to/original_infinigen \
                                     --output_dir /path/to/processed_infinigen_like_hypersim
"""
import argparse
import os
from pathlib import Path
import sys

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm

from utils.cropping import crop_resize_if_necessary
# Ensure OpenEXR support for OpenCV.


def get_parser():
    """Gets the parser for the command-line arguments."""
    parser = argparse.ArgumentParser(
        description="Preprocess the Infinigen dataset into a Hypersim-like format."
    )
    parser.add_argument(
        "--infinigen_dir",
        default='/mnt/mydrive/data/3D/infinigen',
        help="Root directory of the original pre-rendered Infinigen dataset.",
    )
    parser.add_argument(
        "--output_dir",
        default='/mnt/mydrive/data/3D/infinigen_processed',
        help="Root directory for the processed data.",
    )
    parser.add_argument(
        "--img_size",
        type=int,
        nargs=2,
        default=(688, 387),
        help="Target size (width, height) for the processed images.",
    )
    return parser


def process_scene(scene_path, output_root_dir, target_size):
    """
    Processes all frames in a single Infinigen scene directory.

    Args:
        scene_path (Path): Path to the scene folder (e.g., 'seed0').
        output_root_dir (Path): Path to the top-level output directory.
        target_size (tuple): The target (width, height) for resizing.
    """
    print(f"Processing scene: {scene_path.name}")

    # --- Tonemapping parameters from the Hypersim script ---
    gamma = 1.0 / 2.2
    inv_gamma = 1.0 / gamma
    percentile = 90
    brightness_nth_percentile_desired = 0.8

    # Define the base directories for each data type within the scene
    scene_path /= "frames"
    base_image_dir = scene_path / "Image"
    base_depth_dir = scene_path / "Depth"
    base_cam_dir = scene_path / "camview"

    if not (base_image_dir.exists() and base_depth_dir.exists() and base_cam_dir.exists()):
        print(f"  - Base data directories ('Image', 'Depth', 'camview') not found in {scene_path}. Skipping.")
        return

    # Use the camera directories inside "Image" as the reference
    camera_dirs = sorted([d for d in base_image_dir.iterdir() if d.is_dir()])

    if not camera_dirs:
        print(f"  - No camera subdirectories found in {base_image_dir}. Skipping.")
        return

    for camera_dir_path in camera_dirs:
        camera_name = camera_dir_path.name

        # --- Define INPUT paths based on the CORRECTED original Infinigen structure ---
        input_image_dir = base_image_dir / camera_name
        input_depth_dir = base_depth_dir / camera_name
        input_cam_dir = base_cam_dir / camera_name

        if not (input_depth_dir.exists() and input_cam_dir.exists()):
            print(f"  - Skipping camera {camera_name}: Corresponding 'Depth' or 'camview' folder is missing.")
            continue

        # --- Define OUTPUT path for this camera view (Hypersim-like) ---
        output_cam_dir = output_root_dir / scene_path.parent.name / camera_name
        output_cam_dir.mkdir(parents=True, exist_ok=True)

        # Get a sorted list of frame numbers from the input directory
        exr_paths = list(input_image_dir.glob("*.exr"))
        for exr_path in tqdm(exr_paths, desc=f"  - {camera_name}"):
            try:
                # 1. Load Data from original paths
                frame_id = int(exr_path.name.split('_')[-2])
                depth_path = Path(str(exr_path).replace("Image", "Depth").replace('exr', 'npy'))
                cam_path = Path(str(exr_path).replace("Image", "camview").replace('exr', 'npz'))

                if not (exr_path.exists() and depth_path.exists() and cam_path.exists()):
                    print(f'Not all files found for {exr_path}')
                    continue

                color_hdr = cv2.imread(str(exr_path), cv2.IMREAD_UNCHANGED)
                if color_hdr is None:
                    print(f"Warning: Could not read {exr_path}. Skipping.")
                    continue
                if color_hdr.ndim == 3:
                    color_hdr = cv2.cvtColor(color_hdr, cv2.COLOR_BGR2RGB)
                color_hdr = color_hdr.astype(np.float32)

                depth = np.load(depth_path).astype(np.float32)
                depth[~np.isfinite(depth)] = 0.0  # Handle invalid depths
                cam_file = np.load(cam_path)
                intrinsics = cam_file["K"].astype(np.float32)
                pose = cam_file["T"].astype(np.float32)

                # 2. Tonemapping (Hypersim-style)
                brightness = (
                    0.3 * color_hdr[:, :, 0] +
                    0.59 * color_hdr[:, :, 1] +
                    0.11 * color_hdr[:, :, 2]
                )
                brightness_valid = brightness
                eps = 1e-4
                p_current = np.percentile(brightness_valid.astype(np.float32), percentile)
                scale = (np.power(brightness_nth_percentile_desired, inv_gamma) / p_current) if p_current > eps else 0.0

                color_tonemapped = np.power(np.maximum(scale * color_hdr, 0), gamma)
                color_ldr_8bit = (np.clip(color_tonemapped, 0.0, 1.0) * 255).astype(np.uint8)
                img_pil = Image.fromarray(color_ldr_8bit)
                
                # print(color_ldr_8bit.shape, depth.shape, pose.shape, intrinsics.shape, color_hdr.shape, exr_path)
                try:
                    img_pil, depth, intrinsics = crop_resize_if_necessary(
                        img_pil, depth, intrinsics.copy(), target_size, info=frame_id, crop=False
                    )
                except Exception as e:
                    print(f"Error processing {frame_id}: {e}")
                    sys.exit(0)

                # print(img_pil.size, depth.shape)

                # 3. Resizing (No Cropping)
                # original_size = (img_pil.width, img_pil.height)
                # img_resized = img_pil.resize(target_size, Image.Resampling.LANCZOS)
                # depth_resized = cv2.resize(depth, dsize=target_size, interpolation=cv2.INTER_NEAREST)

                # # 4. Adjust Camera Intrinsics
                # intrinsics_resized = intrinsics.copy()
                # sx = target_size[0] / original_size[0]
                # sy = target_size[1] / original_size[1]
                # intrinsics_resized[0, 0] *= sx  # fx
                # intrinsics_resized[1, 1] *= sy  # fy
                # intrinsics_resized[0, 2] *= sx  # cx
                # intrinsics_resized[1, 2] *= sy  # cy

                # 5. Save Processed Data with Hypersim-like naming
                out_rgb_path = output_cam_dir / f"{frame_id:06d}_rgb.png"
                out_depth_path = output_cam_dir / f"{frame_id:06d}_depth.npy"
                out_cam_path = output_cam_dir / f"{frame_id:06d}_cam.npz"

                img_pil.save(out_rgb_path)
                np.save(out_depth_path, depth.astype(np.float32))
                np.savez(
                    out_cam_path,
                    intrinsics=intrinsics.astype(np.float32),
                    pose=pose.astype(np.float32),
                )

            except Exception as e:
                print(f"Error processing frame {frame_id} in {scene_path.name}/{camera_name}: {e}")
                continue
            # break
        # break


def main():
    """Main function to run the preprocessing."""
    parser = get_parser()
    args = parser.parse_args()

    rootdir = Path(args.infinigen_dir)
    outdir = Path(args.output_dir)
    img_size = tuple(args.img_size)

    outdir.mkdir(parents=True, exist_ok=True)

    scenes = sorted([p for p in rootdir.iterdir() if p.is_dir()])
    if not scenes:
        print(f"No scene directories (e.g., 'seed0') found in {rootdir}. Please check the path.")
        return

    for scene_path in scenes:
        process_scene(scene_path, outdir, img_size)
        # break

    print("\nPreprocessing finished.")


if __name__ == "__main__":
    main()
